package rate_limit

import (
	"context"
	"errors"
	"github.com/go-redis/redis_rate/v10"
	"github.com/redis/go-redis/v9"
	"sync"
	"time"
)

type RateLimitConfig struct {
	Key      string
	Rate     int
	TimeUnit time.Duration
}

var limitConfigMap = make(map[string]redis_rate.Limit)

// 全局redis限流器对象
var redisLimiter *redis_rate.Limiter

func RegisterRedisLimiter(redisClient *redis.Client) {
	if redisLimiter == nil {
		redisLimiter = redis_rate.NewLimiter(redisClient)
	}
}

var mutex sync.Mutex

// SetLimitConfig 设置限流配置,最好是写配置文件，动态加载
func SetLimitConfig(config RateLimitConfig) {
	mutex.Lock()
	defer mutex.Unlock()
	switch config.TimeUnit {
	case time.Second:
		limitConfigMap[config.Key] = redis_rate.PerSecond(config.Rate)
	case time.Minute:
		limitConfigMap[config.Key] = redis_rate.PerMinute(config.Rate)
	case time.Hour:
		limitConfigMap[config.Key] = redis_rate.PerHour(config.Rate)
	}
}

// DoWithLimit 在限流条件下执行函数
// ctx: 上下文对象，用于控制请求的取消和超时
// key: 限流器的键，用于区分不同的限流配置
// fn: 需要执行的函数，当请求不被限流时调用
// timeout: 请求等待的最大超时时间，超过这个时间则返回超时错误
// delay: 当请求被限流时的等待时间，在这段时间后再次尝试请求
func DoWithLimit(ctx context.Context, key string, fn func(), timeout time.Duration, delay time.Duration) error {
	start := time.Now()
	for {
		limit, ok := limitConfigMap[key]
		if !ok {
			return errors.New("没有找到相关的限制配置: " + key)
		}
		res, err := redisLimiter.Allow(ctx, key, limit)
		if err != nil {
			return err
		}
		// 如果没有被限流，执行对应函数
		if res.Allowed == 1 {
			fn()
			return nil
		}
		// 如果等待都超时了，就会去告诉他没排队上
		if time.Since(start) > timeout {
			return errors.New("等待超时, key:" + key)
		}
		// 请求被限流，但是没有超时，那就等待一小会
		time.Sleep(delay)
	}
}
